#!/usr/bin/env python3
"""
基于 backtrack_trace.json 绘制事件时间线和事件类型分布图。

示例：
    python scripts/plot_events.py outputs/backtrack_trace.json
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

# === 新增：全局字体设置 ===
# 请根据系统实际安装的字体名称进行调整。
plt.rcParams["font.sans-serif"] = [
    "Noto Sans CJK SC",  # Google Noto 字体
    "SimHei",            # 简黑体（大部分 Linux 发行版需安装）
    "Microsoft YaHei",   # Windows 常见中文字体
    "PingFang SC",       # macOS 常见中文字体之一
    "Arial Unicode MS",  # 其他全 Unicode 字体
    "DejaVu Sans",
]
plt.rcParams["axes.unicode_minus"] = False  # 避免负号显示为方块


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="绘制 AlgViz 回溯事件可视化图表",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "input",
        type=Path,
        help="由 postprocess_ndjson.py 生成的 JSON 文件",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("outputs/plots"),
        help="图表输出目录",
    )
    return parser.parse_args()


def load_data(path: Path) -> Dict:
    with path.open("r", encoding="utf-8") as fh:
        return json.load(fh)


def plot_timeline(events: List[Dict], output_dir: Path) -> None:
    timestamps = [event["timestamp"] for event in events]
    t0 = min(timestamps)
    rel_times = [ts - t0 for ts in timestamps]

    colors = {
        "enter": "#1f77b4",
        "solution": "#2ca02c",
        "pop": "#d62728",
        "default": "#7f7f7f",
    }
    c = [colors.get(event["event"], colors["default"]) for event in events]
    labels = [event["event"] for event in events]

    fig, ax = plt.subplots(figsize=(10, 4))
    ax.scatter(rel_times, [0] * len(events), c=c, s=40)
    ax.set_xlabel("相对时间（秒）")
    ax.set_yticks([])
    ax.set_title("回溯事件时间线")
    ax.grid(axis="x", linestyle="--", alpha=0.3)

    for x, label in zip(rel_times, labels):
        ax.annotate(label, (x, 0), textcoords="offset points", xytext=(0, 8),
                    ha="center", fontsize=8, rotation=90)

    output_dir.mkdir(parents=True, exist_ok=True)
    out_path = output_dir / "timeline.png"
    fig.tight_layout()
    fig.savefig(out_path, dpi=200)
    plt.close(fig)
    print(f"[plot] 时间线保存于：{out_path}")


def plot_event_distribution(stats: Dict, output_dir: Path) -> None:
    events_by_type = stats["events_by_type"]
    items = sorted(events_by_type.items(), key=lambda kv: kv[0])
    labels = [key for key, _ in items]
    counts = [value for _, value in items]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(labels, counts, color="#4c72b0")
    ax.set_title("事件类型分布")
    ax.set_xlabel("事件类型")
    ax.set_ylabel("出现次数")
    ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True))

    output_dir.mkdir(parents=True, exist_ok=True)
    out_path = output_dir / "event_distribution.png"
    fig.tight_layout()
    fig.savefig(out_path, dpi=200)
    plt.close(fig)
    print(f"[plot] 事件分布图保存于：{out_path}")


def main() -> None:
    args = parse_args()
    data = load_data(args.input)

    events = data.get("events", [])
    stats = data.get("stats", {})

    if not events:
        raise ValueError("事件列表为空，请确认输入文件是否正确。")

    plot_timeline(events, args.output)
    if stats:
        plot_event_distribution(stats, args.output)


if __name__ == "__main__":
    main()